/*apr  socket handling */
#include <apr_network_io.h>
#include "util.h"
#include "msg.h"
#include "dtypes.h"
/** \file msg.c
 *  \brief    Return the status of msg structures
 *            REMINDER : TO B FILLED IN DETAIL
 */
apr_status_t msg_recv_type(apr_socket_t *socket, MessageType *type)
{
	apr_size_t bytes_recvd;
	
	return apr_socket_recv_full(socket, type, sizeof(MessageType),
		&bytes_recvd);
}

apr_status_t msg_send_type(apr_socket_t *socket, MessageType type)
{
	apr_size_t bytes_sent;
	
	return apr_socket_send_full(socket, &type, sizeof(MessageType),
		&bytes_sent);
}

apr_status_t msg_recv_env(apr_socket_t *socket, Envelope *env)
{
	apr_size_t bytes_recvd;
	
	return apr_socket_recv_full(socket, env, sizeof(Envelope),
		&bytes_recvd);
}

apr_status_t msg_send_env(apr_socket_t *socket, Envelope *env)
{
	apr_size_t bytes_sent;

	return apr_socket_send_full(socket, env, sizeof(Envelope),
		&bytes_sent);
}

apr_status_t msg_recv_size(apr_socket_t *socket, int *size)
{
	apr_size_t bytes_recvd;
	
	return apr_socket_recv_full(socket, size, sizeof(int), &bytes_recvd);
}

apr_status_t msg_send_size(apr_socket_t *socket, int size)
{
	apr_size_t bytes_sent;
	
	return apr_socket_send_full(socket, &size, sizeof(int), &bytes_sent);
}

int msg_env_match(Envelope *sender, Envelope *receiver)
{
	return
		((sender->source == receiver->source)
			|| (receiver->source == MPI_ANY_SOURCE))
		&&
		(sender->dest == receiver->dest)
		&&
		(sender->context == receiver->context)
		&&
		((sender->tag == receiver->tag)
			|| (receiver->tag == MPI_ANY_TAG));
}

void msg_env_null(Envelope *dst)
{
	dst->source = -1;
	dst->dest = -1;
	dst->context = -1;
	dst->tag = -1;
}

void msg_env_copy(Envelope *dst, Envelope *src)
{
	dst->source = src->source;
	dst->dest = src->dest;
	dst->context = src->context;
	dst->tag = src->tag;
}

static apr_status_t msg_recv_data_basic(apr_socket_t *socket, 
        void *buf, int count, MPI_Datatype datatype, 
        int *total_bytes_recvd)
{
	apr_status_t status;
	int num_bytes, element_size, bytes_remaining, recv_count;
	apr_size_t bytes_recvd;
	char *next_element;

	/* read data length in bytes from socket (unsigned int) */
	
	status = msg_recv_size(socket, &num_bytes);
	if (status != APR_SUCCESS) return status;

	/* determine the size in bytes of each element to be received
	   based on the datatype (currently only works for base datatypes) */

	element_size = sizeof_base_type(datatype);
	if (element_size == 0) return APR_EINVAL;
	
	/* receive elements one at a time until we get the number of
	   elements we are expecting or we run out of bytes */
	
	bytes_remaining = num_bytes;
	recv_count = 0;
	next_element = (char *) buf;
	
	while (recv_count <= count && bytes_remaining >= element_size)
	{
		status = apr_socket_recv_full(socket, next_element, element_size,
			&bytes_recvd);
		bytes_remaining -= bytes_recvd;	
		if (status != APR_SUCCESS || bytes_recvd < element_size) return status;
		recv_count++;
		next_element += element_size;
	}
	
	/* check for leftover bytes in the message; these should be read
	   and discarded to clear them from the socket */
	
	while (bytes_remaining > 0)
	{
		bytes_recvd = 1;
		status = apr_socket_recv(socket, next_element, &bytes_recvd);
		bytes_remaining -= bytes_recvd;	
		if (status != APR_SUCCESS || bytes_recvd < 1) return status;
	}
	
	/* return the total number of bytes received for later use */
	
	*total_bytes_recvd = num_bytes;
	
	return APR_SUCCESS;
}

static apr_status_t msg_recv_data_derived(apr_socket_t *socket, 
        void *buf, int count, MPI_Datatype datatype, 
        int *total_bytes_recvd)
{
	apr_status_t status;
	int num_bytes;
    int position;
	apr_size_t bytes_recvd;
	char *recvbuf;
    MPI_Datatype recvtype;

	/* if count > 1, create a new contiguous type */
    if (count == 1)
        recvtype = datatype;
    else
    {
        MPI_Type_contiguous(count, datatype, &recvtype);
        MPI_Type_commit(&recvtype);
    }

	/* read data length in bytes from socket (unsigned int) */
	status = msg_recv_size(socket, &num_bytes);
	if (status != APR_SUCCESS) return status;

    /*****************************************
     * BEGIN RECEIVE AND UNPACK METHOD       *
     *                                       *
     * Note: This is not an ideal solution.  *
     *      A receive-by-parts method should *
     *      be used in conjunction with a    *
     *      compatible send method.          *
     *****************************************/
    
    /* allocate buffer to receive data */
    recvbuf = (char *)OOM_palloc(num_bytes);
   
    /* receive the data */
    status = apr_socket_recv_full(socket, recvbuf, num_bytes,
			&bytes_recvd);
		
    if (status != APR_SUCCESS || bytes_recvd < num_bytes) return status;
	
    /* unpack the data into the user buffer */
    position = 0;
    MPI_Unpack((void *)recvbuf, num_bytes, &position, buf, 1, 
            recvtype, MPI_COMM_WORLD); /* arbitrary communicator */

	/* return the total number of bytes received for later use */
	*total_bytes_recvd = bytes_recvd;
	
	return APR_SUCCESS;
}

apr_status_t msg_recv_data(apr_socket_t *socket, void *buf, int count,
	MPI_Datatype datatype, int *total_bytes_recvd)
{
    apr_status_t status;
    
    if (datatype < MPI_DATATYPE_BASIC_MAX)
    {
        status = msg_recv_data_basic(socket, buf, count, 
                datatype, total_bytes_recvd);
        return status;
    }
    else
    {
        status = msg_recv_data_derived(socket, buf, count,
                datatype, total_bytes_recvd);
        return status;
    }
}

static apr_status_t msg_send_data_basic(apr_socket_t *socket, 
        void *buf, int count, MPI_Datatype datatype)
{
	apr_status_t status;
	int element_size, num_bytes;
	apr_size_t bytes_sent;

	/* determine the size int bytes of each element to be received
	   based on the datatype (currently only works for base datatypes) */
	
	element_size = sizeof_base_type(datatype);
	if (element_size == 0) return APR_EINVAL;
	
	/* write data length in bytes to socket */

	num_bytes = element_size * count;
	status = msg_send_size(socket, num_bytes);
	if (status != APR_SUCCESS) return status;
	
	/* write the payload to the socket */
	
	status = apr_socket_send_full(socket, (char *) buf, num_bytes,
		&bytes_sent);
	if (status != APR_SUCCESS || bytes_sent < num_bytes) return status;
	
	return APR_SUCCESS;
}

static apr_status_t msg_send_data_derived(apr_socket_t *socket, 
        void *buf, int count, MPI_Datatype datatype)
{
	apr_status_t status;
	int num_bytes, position;
	apr_size_t bytes_sent;
    char *sendbuf;
    MPI_Datatype sendtype;

	/* if count > 1, create a new contiguous type */
    if (count == 1)
        sendtype = datatype;
    else
    {
        MPI_Type_contiguous(count, datatype, &sendtype);
        MPI_Type_commit(&sendtype);
    }
    
    /**********************************************
     * BEING PACK AND SEND METHOD                 *
     *                                            *
     * Note: This is not an ideal final solution. *
     *      A send-by-parts method should be used *
     *      to avoid rebuffering. A compatible    *
     *      receive method is required.           *
     **********************************************/

    /* determine buffer size for packed data
     * NOTE: communicator is arbitrary */
    MPI_Pack_size(1, sendtype, MPI_COMM_WORLD, &num_bytes);
    
    /* allocate buffer */
    sendbuf = (char *)OOM_palloc(num_bytes);
    
    /* pack the data */
    position = 0;
    MPI_Pack(buf, 1, sendtype, (void *)sendbuf, num_bytes, 
            &position, MPI_COMM_WORLD); /* arbitrary communicator */

    /* send packed data */
    status = msg_send_size(socket, num_bytes);
	if (status != APR_SUCCESS) return status;
	
	/* write the payload to the socket */
	status = apr_socket_send_full(socket, sendbuf, num_bytes, 
            &bytes_sent);
	if (status != APR_SUCCESS || bytes_sent < num_bytes) return status;
	
	return APR_SUCCESS;
}

apr_status_t msg_send_data(apr_socket_t *socket, void *buf, int count,
	MPI_Datatype datatype)
{
    apr_status_t status;
    
    if (datatype < MPI_DATATYPE_BASIC_MAX)
    {
        status = msg_send_data_basic(socket, buf, count, datatype);
        return status;
    }
    else
    {
        status = msg_send_data_derived(socket, buf, count, datatype);
        return status;
    }
}

